package com.muyun.data.modeler.infrastructure.general.toolkit;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

import javax.crypto.Cipher;

import com.muyun.data.modeler.infrastructure.general.consts.MessageConsts;
import com.muyun.data.modeler.infrastructure.general.exceptions.TipException;
import org.apache.commons.codec.binary.Base64;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * RSA签名和验签工具类
 *
 * @author 钱多多--付广才
 *
 */
public class RsaToolkit {

    private static final Logger logger = LoggerFactory.getLogger(RsaToolkit.class);

    public static final String RSA = "RSA";
    public static final String MD5_WITH_RSA = "MD5withRSA";
    public static final String SHA1_WITH_RSA = "SHA1WithRSA";
    public static final String SHA256_WITH_RSA = "SHA256WithRSA";
    /**
     * RSA最大加密明文大小
     */
    private static final int MAX_ENCRYPT_BLOCK = 245;
    /**
     * RSA最大解密密文大小
     */
    private static final int MAX_DECRYPT_BLOCK = 256;

    private RsaToolkit() {
        throw new TipException(MessageConsts.CAN_NOT_INSTANCE_CLASS);
    }

    /**
     * 使用MD5WithRSA算法签名数据
     *
     * @param privateKey 私钥
     * @param signByte   待签名的字节
     * @return
     */
    public static byte[] signMD5WithRSA(String privateKey, byte[] signByte) {
        return sign(privateKey, signByte, MD5_WITH_RSA);
    }

    /**
     * 使用SHA1WithRSA算法签名数据(RSA)
     *
     * @param privateKey      私钥
     * @param signByte 待签名的字节
     * @return
     */
    public static byte[] signSHA1WithRSA(String privateKey, byte[] signByte) {
        return sign(privateKey, signByte, SHA1_WITH_RSA);
    }

    /**
     * 使用SHA256WithRSA算法签名数据(RSA2)
     *
     * @param privateKey      私钥
     * @param signByte 待签名的字节
     * @return
     */
    public static byte[] signSHA256WithRSA(String privateKey, byte[] signByte) {
        return sign(privateKey, signByte, SHA256_WITH_RSA);
    }

    /**
     * 使用指定的签名算法签名数据
     *
     * @param privateKey         私钥
     * @param signByte           待签名数据
     * @param signatureAlgorithm 签名算法
     * @return
     */
    public static byte[] sign(String privateKey, byte[] signByte, String signatureAlgorithm) {
        try {
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKey));
            KeyFactory keyf = KeyFactory.getInstance(RSA);
            PrivateKey myprikey = keyf.generatePrivate(keySpec);
            Signature signet = Signature.getInstance(signatureAlgorithm);
            signet.initSign(myprikey);
            signet.update(signByte);
            return signet.sign();
        } catch (Exception e) {
            logger.error("签名异常", e);
        }
        return null;
    }

    /**
     *
     * RSA签名校验验证
     *
     * @param publicKey：RSA公钥
     * @param signByte：待签名源内容
     * @param signedByte：RSA签名结果
     * @return
     */
    public static boolean verifySignSHA1WithRSA(String publicKey, byte[] signByte, byte[] signedByte) {
        return verifySign(publicKey, signByte, signedByte, SHA1_WITH_RSA);
    }

    /**
     *
     * RSA签名校验验证
     *
     * @param publicKey:RSA公钥
     * @param signByte：待签名源内容
     * @param signedByte：RSA签名结果
     * @return
     */
    public static boolean verifySignSHA256WithRSA(String publicKey, byte[] signByte, byte[] signedByte) {
        return verifySign(publicKey, signByte, signedByte, SHA256_WITH_RSA);
    }

    /**
     *
     * RSA签名校验验证
     *
     * @param publicKey：RSA公钥
     * @param signByte：待签名源内容
     * @param signedByte：RSA签名结果
     * @return
     */
    public static boolean verifySignMD5WithRSA(String publicKey, byte[] signByte, byte[] signedByte) {
        return verifySign(publicKey, signByte, signedByte, MD5_WITH_RSA);
    }

    /**
     * 使用指定的签名算法校验签名是否合法
     *
     * @param publicKey          公钥
     * @param signByte           待校验的数据
     * @param signedByte         签名数据
     * @param signatureAlgorithm 签名算法
     * @return
     */
    public static boolean verifySign(String publicKey, byte[] signByte, byte[] signedByte, String signatureAlgorithm) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(RSA);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.decodeBase64(publicKey));
            PublicKey pubKey = keyFactory.generatePublic(keySpec);
            Signature signetcheck = Signature.getInstance(signatureAlgorithm);
            signetcheck.initVerify(pubKey);
            signetcheck.update(signByte);
            return signetcheck.verify(signedByte);
        } catch (Exception e) {
            logger.error("验签异常", e);
        }
        return false;
    }

    /**
     * <P>
     * 私钥解密
     * </p>
     *
     * @param encryptedData 已加密数据
     * @param privateKey    私钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] decryptByPrivateKey(byte[] encryptedData, String privateKey) throws Exception {
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PrivateKey privateK = keyFactory.generatePrivate(pkcs8KeySpec);
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.DECRYPT_MODE, privateK);
        return cipher.doFinal(encryptedData);
    }

    /**
     * <p>
     *    私钥解密，支持任意长度的明文
     * </p>
     * @param encryptedData
     * @param privateKey
     * @return
     * @throws Exception
     */
    // public static byte[] decryptByPrivateKeyWithAnyLength(byte[] encryptedData, String privateKey) throws Exception {
//        byte[] encode = new byte[] {};
//        for (int i = 0; i < bytes.length; i += encodeLen) {
//            byte[] subarray = ArrayUtils.subarray(bytes, i, i + encodeLen);
//            byte[] doFinal = encryptByPublicKey(subarray, key);
//            encode = ArrayUtils.addAll(encode, doFinal);
//        }
//        return encode;
    //  }

    /**
     * <p>
     * 公钥解密
     * </p>
     *
     * @param encryptedData 已加密数据
     * @param publicKey     公钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] decryptByPublicKey(byte[] encryptedData, String publicKey) throws Exception {
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(Base64.decodeBase64(publicKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PublicKey publicK = keyFactory.generatePublic(x509KeySpec);
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.DECRYPT_MODE, publicK);
        return cipher.doFinal(encryptedData);
    }

    /**
     * <p>
     * 公钥解密，支持任意长度，密钥长度必须是2048位
     * </p>
     *
     * @param encryptedData 已加密数据
     * @param publicKey     公钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] decryptByPublicKeyWithAnyLength(byte[] encryptedData, String publicKey) throws Exception {
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(Base64.decodeBase64(publicKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PublicKey publicK = keyFactory.generatePublic(x509KeySpec);
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.DECRYPT_MODE, publicK);
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            byte[] cache;
            int offset = 0;
            int sliceIdx = 0;
            int dataLength = encryptedData.length;
            int remainData = 0;
            while ((remainData = dataLength - offset) > 0) {
                if (remainData > MAX_DECRYPT_BLOCK ) {
                    cache = cipher.doFinal(encryptedData, offset, MAX_DECRYPT_BLOCK);
                } else {
                    cache = cipher.doFinal(encryptedData, offset, remainData);
                }
                out.write(cache, 0, cache.length);
                sliceIdx++;
                offset = sliceIdx * MAX_DECRYPT_BLOCK;
            }
            return out.toByteArray();
        }
    }

    /**
     * <p>
     * 公钥解密，支持任意长度，密钥长度必须是2048位
     * </p>
     *
     * @param encryptedData 已加密数据，会自动进行Base64解码
     * @param publicKey     公钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static String decryptByPublicKeyWithAnyLength(String encryptedData, String publicKey) throws Exception {
        return new String(decryptByPublicKeyWithAnyLength(Base64.decodeBase64(encryptedData), publicKey),
                StandardCharsets.UTF_8);
    }

    /**
     * <p>
     * 公钥加密
     * </p>
     *
     * @param data      源数据
     * @param publicKey 公钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] encryptByPublicKey(byte[] data, String publicKey) throws Exception {
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(Base64.decodeBase64(publicKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PublicKey publicK = keyFactory.generatePublic(x509KeySpec);
        // 对数据加密
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, publicK);
        return cipher.doFinal(data);
    }

    /**
     * <p>
     * 私钥加密
     * </p>
     *
     * @param data       源数据
     * @param privateKey 私钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] encryptByPrivateKey(byte[] data, String privateKey) throws Exception {
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PrivateKey privateK = keyFactory.generatePrivate(pkcs8KeySpec);
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, privateK);
        return cipher.doFinal(data);
    }

    /**
     * <p>
     * 私钥加密
     * </p>
     *
     * @param data       源数据
     * @param privateKey 私钥(BASE64编码)
     * @return Base64编码的密文
     * @throws Exception
     */
    public static String encryptBase64ByPrivateKey(byte[] data, String privateKey) throws Exception {
        return Base64.encodeBase64String(encryptByPrivateKey(data, privateKey));
    }

    /**
     * <p>
     * 私钥加密
     * </p>
     *
     * @param data       源数据
     * @param privateKey 私钥(BASE64编码)
     * @return Base64编码的密文
     * @throws Exception
     */
    public static String encryptBase64ByPrivateKey(String data, String privateKey) throws Exception {
        return Base64.encodeBase64String(encryptByPrivateKey(data.getBytes(StandardCharsets.UTF_8), privateKey));
    }

    /**
     * <p>
     * 私钥加密
     * </p>
     *
     * @param data       源数据
     * @param privateKey 私钥(BASE64编码)
     * @return Base64编码的密文
     * @throws Exception
     */
    public static String encryptBase64ByPrivateKey(String data, String privateKey, String charset) throws Exception {
        return Base64.encodeBase64String(encryptByPrivateKey(data.getBytes(charset), privateKey));
    }

    /**
     * <p>
     * 私钥加密，支持任意长度，密钥必须位2048位
     * </p>
     *
     * @param data       源数据
     * @param privateKey 私钥(BASE64编码)
     * @return
     * @throws Exception
     */
    public static byte[] encryptByPrivateKeyWithAnyLength(byte[] data, String privateKey) throws Exception {
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        PrivateKey privateK = keyFactory.generatePrivate(pkcs8KeySpec);
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, privateK);
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            byte[] cache;
            int offset = 0;
            int sliceIdx = 0;
            int dataLength = data.length;
            int remainData = 0;
            while ((remainData = dataLength - offset) > 0) {
                if (remainData > MAX_ENCRYPT_BLOCK ) {
                    cache = cipher.doFinal(data, offset, MAX_ENCRYPT_BLOCK);
                } else {
                    cache = cipher.doFinal(data, offset, remainData);
                }
                out.write(cache, 0, cache.length);
                sliceIdx++;
                offset = sliceIdx * MAX_ENCRYPT_BLOCK;
            }
            return out.toByteArray();
        }
    }

    public static Map<String, String> genKey() throws NoSuchAlgorithmException {
        Map<String, String> keyMap = new HashMap<String, String>();
        KeyPairGenerator keygen = KeyPairGenerator.getInstance(RSA);
        SecureRandom random = new SecureRandom();
        // random.setSeed(keyInfo.getBytes());
        // 初始加密，512位已被破解，用1024位,最好用2048位
        keygen.initialize(2048, random);
        // 取得密钥对
        KeyPair kp = keygen.generateKeyPair();
        RSAPrivateKey privateKey = (RSAPrivateKey) kp.getPrivate();
        String privateKeyString = Base64.encodeBase64String(privateKey.getEncoded());
        RSAPublicKey publicKey = (RSAPublicKey) kp.getPublic();
        String publicKeyString = Base64.encodeBase64String(publicKey.getEncoded());
        keyMap.put("publicKey", publicKeyString);
        keyMap.put("privateKey", privateKeyString);
        System.out.println(keyMap.get("publicKey"));
        System.out.println(keyMap.get("privateKey"));
        return keyMap;
    }

}
